import csv
import datetime
import queue
import os
import sys
import time
import ssl

from paho.mqtt import client as mqtt_client


'''
    从db中查询有效可用的所有测试车辆信息数据
    根据命令行的输入参数，指定启动的车辆数，以及与broker代理建立连接的频率
    建立连接成功的车辆，就可以根据脚本里指定的频次，来像broker发送数据
    脚本统计连接数、请求数、响应时间等信息写到报表中
    调试遇到车辆会批量断开连接的情况，增加了当车辆断开连接时，把断开时间、车辆信息写到本地csv中，方便第二天来查看分析。
'''

# 根据不同系统进行路径适配
if os.name == "nt":
    path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
    sys.path.insert(0, path)
    from GB_test.utils.mysql_operating import DB
elif os.name == "posix":
    sys.path.append("/app/qa_test_app/")
    from GB_test.utils.mysql_operating import DB

from locust import User, TaskSet, events, task, between, run_single_user


BROKER_ADDRESS = "broker服务地址"
PORT = 1111
PASSWORD = "111111"
PUBLISH_TIMEOUT = 10000  # 超时时间
TEST_TOPIC = "test_topic"

TEST_VALUE = [16, 3, -26, 4, 0, 36]  # 用来publish的测试数据，仅示意

BYTES_DATA = bytes(i % 256 for i in TEST_VALUE)  # 业务需要转换成 byte 类型后再发送

# 创建队列
client_queue = queue.Queue()

# 连接DB，读取车辆数据
db = DB("db_vmd")
select_sql = "select xxxx"
client_list = db.fetch_all(select_sql)
print("车辆数据查询完毕，数据量：{}".format(len(client_list)))
for t in client_list:
    # 把可用的车辆信息存到队列中去
    client_queue.put(t)


def fire_success(**kwargs):
    """请求成功时调用"""
    events.request.fire(**kwargs)


def calculate_resp_time(t1, t2):
    """计算响应时间"""
    return int((t2 - t1) * 1000)


class MQTTMessage:
    """已发送的消息实体类"""
    def __init__(self, _type, qos, topic, payload, start_time, timeout):
        self.type = _type,
        self.qos = qos,
        self.topic = topic
        self.payload = payload
        self.start_time = start_time
        self.timeout = timeout


# 统计总共发送成功的消息数量
total_published = 0
disconnect_record_list = []  # 定义存放连接断开的记录的列表容器


class PublishTask(TaskSet):

    @task
    def task_publish(self):
        self.client.loop_start()
        topic = TEST_TOPIC
        payload = BYTES_DATA
        # 记录发送的开始时间
        start_time = time.time()
        mqtt_msg_info = self.client.publish(topic, payload, qos=1, retain=False)
        published_mid = mqtt_msg_info.mid
        # 将发送成功的消息内容，放入client实例的 published_message 字段
        self.client.published_message[published_mid] = MQTTMessage(REQUEST_TYPE,
                                                                   0,
                                                                   topic,
                                                                   payload,
                                                                   start_time,
                                                                   PUBLISH_TIMEOUT)
        # 发送成功回调
        self.client.on_publish = self.on_publish
        # 断开连接回调
        self.client.on_disconnect = self.on_disconnect

    @staticmethod
    def on_disconnect(client, userdata, rc):
        """ broker连接断开，放入列表容器"""
        disconnected_info = [str(client._client_id), rc, datetime.datetime.now()]
        disconnect_record_list.append(disconnected_info)
        print("rc状态:{} - -".format(rc), "{}-broker连接已断开".format(str(client._client_id)))

    @staticmethod
    def on_publish(client, userdata, mid):
        if mid:
            # 记录消息发送成功的时间
            end_time = time.time()
            # 从已发送的消息容器中，取出消息
            message = client.published_message.pop(mid, None)
            # 计算开始发送到发送成功的耗时
            publish_resp_time = calculate_resp_time(message.start_time, end_time)
            fire_success(
                request_type="p_success",
                name="client_id: " + str(client._client_id),
                response_time=publish_resp_time,
                response_length=len(message.payload),
                exception=None,
                context=None
            )
            global total_published
            # 成功发送累加1
            total_published += 1


class MQTTLocustUser(User):
    tasks = [PublishTask]
    wait_time = between(2, 2)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 从队列中获取客户端 username 和 client_id
        current_client = client_queue.get()

        self.client = mqtt_client.Client(current_client[1])
        self.client.username_pw_set(current_client[0], PASSWORD)
        # self.client.username_pw_set(current_client[0] + "1", PASSWORD)  # 模拟client连接报错

        # 定义一个容器，存放已发送的消息
        self.client.published_message = {}

    def on_start(self):
        # 设置tls
        context = ssl.SSLContext(ssl.PROTOCOL_TLS)
        self.client.tls_set_context(context)

        self.client.connect(host=BROKER_ADDRESS, port=PORT, keepalive=60)
        self.client.on_connect = self.on_connect

    def on_stop(self):
        print("publish 成功, 当前已成功发送数量：{}".format(total_published))
        if len(disconnect_record_list) == 0:
            print("无断开连接的client")
        else:
            # 把断开记录里的信息写入csv
            with open("disconnect_record.csv", "w", newline='', encoding='UTF8') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(['client_id', 'rc_status', 'disconnected_time'])
                for i in disconnect_record_list:
                    writer.writerow(i)
            print("断开连接的client信息已写入csv文件")

    @staticmethod
    def on_connect(client, userdata, flags, rc, props=None):
        if rc == 0:
            print("rc状态:{} - -".format(rc), "{}-连接broker成功".format(str(client._client_id)))
            fire_success(
                request_type="c_success",
                name='count_connected',
                response_time=0,
                response_length=0,
                exception=None,
                context=None
            )
        else:
            print("rc状态:{} - -".format(rc), "{}-连接broker失败".format(str(client._client_id)))
            fire_success(
                request_type="c_fail",
                name="client_id: " + str(client._client_id),
                response_time=0,
                response_length=0,
                exception=None,
                context=None
            )


if __name__ == '__main__':
    run_single_user(MQTTLocustUser)
